import pandas as pd

pd_all = pd.read_csv('data/ChnSentiCorp_htl_all.csv', index_col=None)
# 还是要处理缺失值
# <class 'float'>
# Traceback (most recent call last):
#   File "C:\Users\Add阿咚东\PycharmProjects\SST-2\tokenProcess.py", line 8, in <module>
#     tokenizer(t, truncation=True, padding='max_length', max_length=512)
#   File "C:\Users\Add阿咚东\AppData\Local\Programs\Python\Python39\lib\site-packages\transformers\tokenization_utils_base.py", line 2356, in __call__
#     raise ValueError(
# ValueError: text input must of type `str` (single example), `List[str]` (batch or single pretokenized example) or `List[List[str]]` (batch of pretokenized examples).

pd_all = pd_all.dropna()
print(pd_all)

# 评论数目（总体）：7766
# 评论数目（正向）：5322
# 评论数目（负向）：2444
# 平衡一下，又尽可能多地拿数据

pd_positive = pd_all[pd_all.label == 1]
pd_negative = pd_all[pd_all.label == 0]


def get_balance_corpus(corpus_size, corpus_pos, corpus_neg):
    sample_size = corpus_size // 2
    pd_corpus_balance = pd.concat([corpus_pos.sample(sample_size, replace=corpus_pos.shape[0] < sample_size),
                                   corpus_neg.sample(sample_size, replace=corpus_neg.shape[0] < sample_size)])

    print('评论数目（总体）：%d' % pd_corpus_balance.shape[0])
    print('评论数目（正向）：%d' % pd_corpus_balance[pd_corpus_balance.label == 1].shape[0])
    print('评论数目（负向）：%d' % pd_corpus_balance[pd_corpus_balance.label == 0].shape[0])

    return pd_corpus_balance


# 先做一次全部负样本抽样，然后正样本取2444个出来,负样本取2444个出来
ChnSentiCorp_htl_ba_2444 = get_balance_corpus(2444*2, pd_positive, pd_negative)

print(ChnSentiCorp_htl_ba_2444)

# 固定每次训练的样本，控制变量
ChnSentiCorp_htl_ba_2444.to_csv('data/ChnSentiCorp_htl_ba_2444.csv')
